import json

from fastapi import (
    APIRouter,
    Depends,
    HTTPException,
    Query,
    WebSocket,
    WebSocketDisconnect,
)
from jose import JWTError, jwt

from app.api.deps import CurrentUserDep
from app.core import settings
from app.schemas.stomp import StompFrame
from app.websocket.connection_manager import manager
from app.websocket.stomp_handler import handle_stomp_message

router = APIRouter(prefix="/websocket", tags=["07.Websocket接口"])


async def get_token_from_websocket(websocket: WebSocket, token: str = Query(...)):
    """
    从 WebSocket URL 查询参数中提取并验证 JWT token
    """
    if not token:
        raise HTTPException(status_code=401, detail="Missing token")

    try:
        payload = jwt.decode(token, settings.SECRET_KEY, algorithms=[settings.ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise JWTError()
        return username
    except JWTError:
        raise HTTPException(status_code=401, detail="Invalid or expired token")


@router.websocket("/websocket/ws")
async def websocket_endpoint(websocket: WebSocket, user_id: str = Depends(get_token_from_websocket)):
    # 握手阶段已认证，user_id 有效
    await websocket.accept()

    # 后续逻辑不变
    await manager.connect(websocket, user_id)
    online_count = len(manager.active_connections)
    await manager.broadcast_to_topic("online-count", str(online_count))

    try:
        while True:
            data = await websocket.receive_text()
            frame = StompFrame.parse(data)
            response = await handle_stomp_message(user_id, frame)
            if response:
                await websocket.send_text(response)
    except WebSocketDisconnect:
        pass
    finally:
        await manager.disconnect(websocket, user_id)
        online_count = len(manager.active_connections)
        await manager.broadcast_to_topic("online-count", str(online_count))


# 用于后端主动触发字典更新（例如管理端修改字典后调用）
@router.post("/api/dict/notify-update")
async def notify_dict_update(dict_code: str, current_user: CurrentUserDep):
    import time

    message = json.dumps({"dictCode": dict_code, "timestamp": int(time.time() * 1000)})
    await manager.broadcast_to_topic("dict", message)
    return {"status": "notified"}
